# ==============================================================================
# 文件: em_trainer.py
# 作用: 实现算法的第二阶段。接收一个训练好的初始模型和数据嵌入，
#      通过EM算法进行迭代优化，利用所有未标记数据来修正模型偏差，
#      并返回最终的分类器。
# ==============================================================================
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset
from tqdm import tqdm
import copy # 导入copy库，用于深度拷贝模型参数

import config # 从同级目录导入配置
from models import SimpleMLP # 从同级目录的models模块导入MLP模型

# 定义加权BCE损失函数
def weighted_bce_with_logits(outputs, targets, weights_p, weights_n):
    """
    自定义的加权BCE损失函数。
    :param outputs: 模型的logit输出。
    :param targets: 真实标签（这里主要用于P集）。
    :param weights_p: 样本为正例的权重（对于U集是软标签w，对于P集是1）。
    :param weights_n: 样本为负例的权重（对于U集是1-w，对于P集是0）。
    :return: 计算出的加权损失。
    """
    # 使用BCEWithLogitsLoss计算每个样本的损失，不进行任何聚合
    loss_per_sample = nn.BCEWithLogitsLoss(reduction='none')(outputs, targets)
    
    # 根据权重调整损失
    # 正例的损失项 * 正例权重 + 负例的损失项 * 负例权重
    # 注意: BCE(output, 1) = log(1 + exp(-output)) -> 正例损失
    #       BCE(output, 0) = log(1 + exp(output))  -> 负例损失
    # 这里我们用一个更简洁的实现：
    # L = w_p * y * log(f) + w_n * (1-y) * log(1-f)
    # 对于P集, y=1, w_p=1, w_n=0, loss = log(f)
    # 对于U集, y=w (软标签), w_p=w, w_n=1-w, loss = w*log(f) + (1-w)*log(1-f)
    # BCEWithLogitsLoss已经帮我们处理了log-sum-exp的数值稳定性
    # 我们只需要对最终的loss_per_sample进行加权
    # 这里targets是0或1的硬标签，我们用它来区分正负损失项
    weighted_loss = weights_p * targets * loss_per_sample + weights_n * (1 - targets) * loss_per_sample
    
    return weighted_loss.mean() # 返回加权损失的均值

def train_final_model_with_em(initial_model, p_embeddings, u_embeddings):
    """
    使用EM算法训练最终分类器。
    :param initial_model: 训练好的初始模型。
    :param p_embeddings: Tensor, P集的嵌入向量。
    :param u_embeddings: Tensor, U集的嵌入向量。
    :return: 最终训练好的MLP模型。
    """
    print("\n--- 阶段二：开始使用EM算法优化模型 ---")
    
    # 初始化模型和优化器
    model = initial_model # 从初始模型开始
    optimizer = torch.optim.Adam(model.parameters(), lr=config.EM_TRAIN_LR)

    # 准备P集和U集的数据加载器
    # P集在M步中用于计算有监督损失
    p_dataset = TensorDataset(p_embeddings, torch.ones(p_embeddings.shape[0], 1))
    p_dataloader = DataLoader(p_dataset, batch_size=config.EM_TRAIN_BATCH_SIZE, shuffle=True)
    
    # U集用于计算加权损失
    u_dataset = TensorDataset(u_embeddings)
    u_dataloader = DataLoader(u_dataset, batch_size=config.EM_TRAIN_BATCH_SIZE, shuffle=False)
    
    # EM迭代循环
    for iteration in range(config.EM_MAX_ITERATIONS):
        print(f"\nEM 迭代 [{iteration+1}/{config.EM_MAX_ITERATIONS}]")
        
        # --- E-步: 计算软标签 (后验概率) ---
        print("E-步: 正在为U集计算软标签...")
        model.eval() # 设置为评估模式
        all_weights = []
        with torch.no_grad(): # 禁用梯度计算
            for batch_u_embeds, in u_dataloader:
                batch_u_embeds = batch_u_embeds.to(config.DEVICE)
                # 前向传播得到logits，然后通过sigmoid得到概率
                outputs = model(batch_u_embeds)
                weights = torch.sigmoid(outputs)
                all_weights.append(weights.cpu())
        
        # 将所有U集的软标签拼接起来
        u_weights = torch.cat(all_weights, dim=0)

        # --- 新增：在每次E步后，打印当前的估计比例 ---
        current_estimated_prior = u_weights.mean().item()
        print(f"E-步完成。当前对U集的正例比例估计 (pi_hat): {current_estimated_prior:.4f}")

        # --- M-步: 更新模型参数 ---
        print("M-步: 正在使用加权损失更新模型...")
        model.train() # 设置为训练模式
        
        # 记录旧的参数以备检查收敛
        old_params = copy.deepcopy(list(model.parameters()))

        # --- 新的M步训练逻辑: 将P和U合并进行训练 ---
        # 1. 创建P集的目标标签（硬标签1）
        p_labels = torch.ones(p_embeddings.shape[0], 1)
        # 2. 将P集和U集的数据和目标（硬/软标签）合并
        m_step_embeddings = torch.cat([p_embeddings, u_embeddings], dim=0)
        m_step_labels = torch.cat([p_labels, u_weights], dim=0) # P集用硬标签，U集用软标签
        # 3. 创建统一的数据集和加载器
        m_step_dataset = TensorDataset(m_step_embeddings, m_step_labels)
        m_step_dataloader = DataLoader(m_step_dataset, batch_size=config.EM_TRAIN_BATCH_SIZE, shuffle=True)
        
        # 在每个M步中训练数个轮次
        for epoch in range(config.EM_TRAIN_EPOCHS_PER_M_STEP):
            for batch_embeds, batch_labels in m_step_dataloader:
                batch_embeds = batch_embeds.to(config.DEVICE)
                batch_labels = batch_labels.to(config.DEVICE)

                optimizer.zero_grad()
                outputs = model(batch_embeds)
                
                # 直接将混合了硬标签和软标签的batch_labels作为目标
                loss = nn.BCEWithLogitsLoss()(outputs, batch_labels)
                
                loss.backward()
                optimizer.step()

        # --- 检查收敛 ---
        param_change = 0
        new_params = list(model.parameters())
        for old_p, new_p in zip(old_params, new_params):
            param_change += torch.norm(old_p.data - new_p.data).item()
        
        print(f"M-步完成。参数变化量: {param_change:.6f}")
        if param_change < config.EM_CONVERGENCE_EPSILON:
            print("EM算法已收敛。")
            break
            
    # 保存最终模型
    print(f"EM算法完成，保存最终模型到: {config.FINAL_MODEL_PATH}")
    torch.save(model.state_dict(), config.FINAL_MODEL_PATH)
    
    return model
